import numpy as np
import matplotlib.pyplot as plt
import cv2 as cv

feature_params = dict(maxCorners=300, qualityLevel=0.2,
                      minDistance=2, blockSize=7)

# Parameters for Lucas-Kanade optical flow
lk_params = dict(winSize=(15, 15), maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT,
                           10, 0.03))

# The video feed is read in as a VideoCapture object
file = "C:/Users/18873/Desktop/videos/shibuya.mp4"
# cap = cv.VideoCapture("C:/Users/18873/Desktop/videos/v_Archery_g20_c06.avi")
cap = cv.VideoCapture(file)
# Variable for color to draw optical flow track
color = (0, 255, 0)
ret, first_frame = cap.read()
prev_gray = cv.cvtColor(first_frame, cv.COLOR_BGR2GRAY)
prev = cv.goodFeaturesToTrack(prev_gray, mask=None, **feature_params)
# Creates an image filled with zeros intensities with the
# same dimensions as the frame - for later drawing purpose
mask = np.zeros_like(first_frame)

while cap.isOpened():
    ret, frame = cap.read()
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    # gray = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
    next, status, error = cv.calcOpticalFlowPyrLK(prev_gray, gray, prev, None, **lk_params)
    # Selects good feature points for previous position
    good_old = prev[status == 1]
    # Selects good feature points for next position
    good_new = next[status == 1]
    # Draws the optical flow tracks
    for i, (new, old) in enumerate(zip(good_new, good_old)):
        # Returns a contiguous flattened
        # array as (x,y) coordinates for new point
        a, b = new.ravel()
        # Returns a contiguous flattened array as
        # (x,y) coordinates for old point
        c, d = old.ravel()
        # Draws line between new and old position with
        # green color and 2 thickness
        mask = cv.line(mask, (a, b), (c, d), color, 2)
        frame = cv.circle(frame, (a, b), 3, color, -1)
    output = cv.add(frame, mask)
    prev_gray  = gray.copy()
    prev = good_new.reshape(-1, 1, 2)
    cv.imshow('sparse optical flow', output)

    if cv.waitKey(10) & 0xFF == ord('q'):
        break

# The following frees up resources and closes all windows
cap.release()
cv.destroyAllWindows()
